from __future__ import annotations
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams

import _config as cfg
import _gkd as gkd


# 作图默认参数控制
rcParams['savefig.dpi'] = 300  # 图片保存默认分辨率
rcParams['savefig.transparent'] = True  # 图片保存透明背景
rcParams['font.size'] = 12  # 图片默认字号
rcParams['axes.unicode_minus'] = False  # 作图时正常显示符号
rcParams['axes.linewidth'] = 1.2  # spine 边框线宽
rcParams['font.sans-serif'] = 'Microsoft YaHei'  # 作图字体
for v in ['xtick', 'ytick']:
    rcParams[v + '.major.size'] = 7  # 主刻度线长
    rcParams[v + '.minor.size'] = 4  # 次刻度线长
    rcParams[v + '.major.width'] = 1.2  # 主刻度线宽
    rcParams[v + '.minor.width'] = 1.2  # 次刻度线宽


def plot_annual_by_species(data_: dict, window_title=''):
    """ 年均值作图, 子图对应物种 
        data_: dict, {'PM2.5': df1, 'O3': df2}
                df1: index为DateTimeIndex, columns为省市站点名称

        window_title: 作图窗口标题

    无返回值
    2023-09-04
    """

    # 物种列表
    list_species = list(data_.keys())
    
    # 省/市/站点列表
    list_regeion = data_[list_species[0]].columns.tolist()

    plt.ion()

    # 画布大小
    canvas_width, canvas_height = [int(i / cfg.dpi) for i in cfg.canvas_size]

    # 颜色控制
    if len(list_regeion) <= 10:
        colors = plt.cm.tab10(np.linspace(0, 1, 10))  # y值数量10个以下时用tab10颜色
    else:
        colors = plt.cm.jet(np.linspace(0, 0.9, len(list_regeion)))  # 10个以上用jet颜色colormap

    # 作图
    rows = round(len(list_species) ** 0.5)  # 作图行数
    cols = math.ceil(len(list_species) / rows)  # 作图列数

    fig, ax = plt.subplots(rows, cols, figsize=(canvas_width, canvas_height), dpi=cfg.dpi)

    if isinstance(ax, np.ndarray):
        ax = ax.flatten()
    else:
        ax = [ax]

    total_width = 0.8
    bar_width = total_width / len(list_regeion)
    # loc_x0 = [0.5 * (bar_width - total_width) + i for i in range(4)]

    for k in range(len(list_species)):
        df_k = data_[list_species[k]]

        list_x = df_k.index.year

        loc_x0 = [0.5 * (bar_width - total_width) + i for i in range(len(list_x))]

        for n in range(len(list_regeion)):
            loc_x = [i + n * bar_width for i in loc_x0]
            values = df_k.iloc[:, n].values
            ax[k].bar(loc_x, values, width=bar_width, color=colors[n])

        # 设置x轴刻度
        ax[k].set_xticks(np.arange(len(list_x)))

        # x轴刻度label旋转
        ax[k].set_xticklabels(list_x, rotation=45, ha='center')
        # ax[k].set_xticklabels(list_x, rotation=45, ha='right')

        # 子图标题
        ax[k].set_title(list_species[k])

    # 设置公共legend
    plt.gcf().legend(
        # ax[0] + ax[1],
        # ['spring', 'summer', 'autumn', 'winter'],
        data_[list_species[0]].columns,
        # loc='upper left',
        loc='lower center',
        bbox_to_anchor=(0.5, 0.945),
        ncol=data_[list_species[0]].shape[1],
        frameon=False,
        # fontsize=16,
        # labelspacing=0,
        handlelength=1,  # 图例的长度
        handletextpad=0.1,  # 图例与文字的间距
        columnspacing=0.2,  # 图例列间距
    )

    # 窗口标题
    fig.canvas.manager.set_window_title(window_title)
    # fig.canvas.manager.set_window_title('Plot_annual_by_species')

    plt.tight_layout()
    plt.subplots_adjust(top=0.92)
    plt.show()


def plot_annual_by_region(data_: dict, window_title=''):
    """ 年均值作图, 子图对应区域
        data_: dict, {'PM2.5': df1, 'O3': df2}
                df1: index为DateTimeIndex, columns为省市站点名称
        
        window_title: 作图窗口标题
    
    无返回值
    2023-09-04
    """

    # 判断CO是否在其中
    if 'CO' in data_.keys():
        data_['CO×100'] = data_['CO'] * 100
        
        # 删除CO
        del data_['CO']

    # 物种列表
    list_species = list(data_.keys())

    
    # 省/市/站点列表
    list_regeions = data_[list_species[0]].columns.tolist()

    # 数据转换
    dict_data = dict()
    for r in list_regeions:

        # 合并数据
        df_r = pd.concat(objs=[data_[s].loc[:, r] for s in list_species], axis=1)
        
        # 设置表头
        df_r.columns = list_species

        # 数据添加至字典
        dict_data[r] = df_r

    # 作图
    plot_annual_by_species(data_=dict_data, window_title=window_title)


def plot_primary_pollutant(data: pd.DataFrame, window_title=''):
    """ 首要污染物统计

        data_: index为月分辨率DateTimeIndex, columns为['O3', 'SO2', 'NO2', 'CO', 'PM10', 'PM2.5']
                值为占比

    无返回值
    2023-09-04
    """

    plt.ion()

    # 提取数据
    df_count = data

    # 物种顺序
    list_species = ['O3', 'SO2', 'NO2', 'CO', 'PM10', 'PM2.5', 'AQI≤50']

    # 表头顺序调整
    list_header = [i for i in list_species if i in df_count.columns]
    df_count = df_count.loc[:, list_header]

    # 年份列表
    list_year = df_count.index.year
    list_year = sorted(list(set(list_year)), reverse=False)
    # # print(list_year)

    """ 作图 """
    # 画布
    fig, ax = plt.subplots(nrows=1, ncols=len(list_year), figsize=(25, 5), dpi=100, sharey=True)
    if len(list_year) > 1:
        ax = ax.flatten()
    else:
        ax = [ax]

    # 颜色列表
    # list_color = ['#7ac70c', '#faa918', '#d33131', '#1cb0f6', '#8549ba', '#4c4c4c']
    dict_color = dict(zip(list_species, cfg.list_color_primary_pollutant))

    # 按年画图
    for i in range(len(list_year)):
        # 提取当前年的数据
        df_i = df_count[df_count.index.year == list_year[i]]

        # 新的索引
        new_index = pd.date_range(start=df_i.index[0].strftime('%Y') + '-01-01', end=df_i.index[-1].strftime('%Y-%m-%d'), freq='MS')

        # 填充月份
        # new_index = pd.date_range(start='2023-01-01', end='2023-07-01', freq='MS')

        # 使用新时间索引重新索引 DataFrame，将会在 DataFrame 中创建 NaN 值
        df_i = df_i.reindex(new_index)

        # 指定月份为索引
        df_i.index = df_i.index.month

        # 画图
        df_i.plot.bar(ax=ax[i], stacked=True, legend=False, width=1, color=[dict_color[k] for k in df_count.columns])

        ax[i].set_xlim(-0.5, 11.5)
        ax[i].set_ylim(0, 100)

        ax[i].set_xlabel(list_year[i])

        # x轴旋转
        ax[i].tick_params(axis='x', which='major', labelrotation=0)

        # 隐藏中间的右边框
        if i < len(list_year) - 1:
            ax[i].spines['right'].set_visible(False)

        if i > 0:
            # ax[i].spines['left'].set_visible(False)
            ax[i].yaxis.set_tick_params(size=0)

    # 图例文字
    list_columns = [cfg.dict_sp2ecies_latex[s] if s in cfg.dict_sp2ecies_latex.keys() else s for s in df_count.columns]

    # 图例
    fig.legend(
        list_columns,
        # df_count.columns,
        loc='upper left',
        # loc='lower center',
        bbox_to_anchor=(0.03, 1.01),
        ncol=df_count.shape[1],
        frameon=False,
        # fontsize=16,
        # labelspacing=0,
        # colors='jet',
        handlelength=2,  # 图例的长度
        handletextpad=0.5,  # 图例与文字的间距
        columnspacing=0.5,  # 图例列间距
    )

    # 窗口标题
    fig.canvas.manager.set_window_title(window_title)

    # y轴label
    ax[0].set_ylabel('Percentage of each primary pollutant (%)')

    plt.tight_layout()
    plt.subplots_adjust(wspace=0, top=0.93)
    plt.show()


if __name__ == '__main__':
    pass
